# -*- coding: utf-8 -*-
'''
@File    :   aliLLM.py
@Author  :   一力辉 
'''

from ..builder import LLMEngines
from ..engineBase import BaseEngine
import json
from dashscope import Application
from http import HTTPStatus
from typing import List, Optional
from digitalHuman.utils import TextMessage
from digitalHuman.utils import logger

__all__ = ["AliAPI"]

@LLMEngines.register("AliAPI")
class AliAPI(BaseEngine): 
    def checkKeys(self) -> List[str]:
        return ["KEY", "APPID"]
    
    def setup(self):
        pass
    
    async def run(self, input: TextMessage, **kwargs) -> Optional[TextMessage]:
        try: 
            responses = Application.call(app_id=self.cfg.APPID,#YOUR-APPID
                                 prompt=input.data,
                                 api_key=self.cfg.KEY, #YOUR-API-KEY
                                 stream=True,
                                 incremental_output=True
                                 )

            for response in responses:
                if response.status_code != HTTPStatus.OK:
                    print('request_id=%s, code=%s, message=%s\n' % (
                        response.request_id, response.status_code, response.message))            
                else:
                    print('output=%s, usage=%s\n' % (response.output.text, response.usage))
                    message = TextMessage(data=response.output.text)
                    yield message
        except Exception as e:
            logger.error(f"[LLM] Engine run failed: {e}")
            yield None